import numpy as np

def topk_np(arr,k,dim):
    idx=np.argpartition(-arr,kth=k,axis=dim)
    idx=idx.take(indices=range(k),axis=dim)
    val=np.take_along_axis(arr,indices=idx,axis=dim)
    sorted_idx=np.argsort(-val,axis=dim)
    idx=np.take_along_axis(idx,indices=sorted_idx,axis=dim)
    val=np.take_along_axis(val,indices=sorted_idx,axis=dim)
    return val,idx


if __name__ == "__main__":
    a = np.random.randn(10)
    score, idx = topk_np(a, 5, dim=-1)
    print(a.tolist())
    print(idx.tolist())
    print(score.tolist())